from yolov11trt.config import *
import ctypes
import sys
from yolov11trt.test_YoloV11TRT import *
from cv.Camera import *
from cv.video import *
from pan_tilt.test import *

from PyQt5.QtWidgets import QApplication, QWidget, QLabel, QSlider, QPushButton, QVBoxLayout, QHBoxLayout, QSizePolicy
from PyQt5.QtCore import Qt, QThread, pyqtSignal
from PyQt5.QtGui import QImage, QPixmap


class App(QWidget):
    def __init__(self, yolov11_wrapper, cam, yuntai):
        super().__init__()
        self.setWindowTitle("PyQt YOLOv11 Trace")
        
        self.yolov11_wrapper = yolov11_wrapper
        self.disply_width = 700
        self.display_height = 500
        self.load_model_status = False
        
        self.resize(self.disply_width, self.display_height)

        self.image_label = QLabel(self)
        self.image_label.resize(640, 480)
        self.cam_closed_pixmap = QPixmap("source/cam_closed.png")
        self.image_label.setPixmap(self.cam_closed_pixmap)
        
        
        self.video_thread = VideoThread(yolov11_wrapper, cam)
        self.video_thread.change_pixmap_signal.connect(self.update_image)

        self.yuntai = yuntai
        self.yuntai_thread = YuntaiThread(yuntai, 640, 480)
        self.video_thread.infer_results_signal.connect(self.yuntai_thread.handle_infer_result)
        self.yuntai_thread.start()

        model_h_layout = QHBoxLayout()
        self.load_model_Button = QPushButton("加载模型", self)
        self.load_model_Button.clicked.connect(self.load_model)
        self.infer_button = QPushButton("开启推理", self)
        self.infer_button.setDisabled(True)
        self.infer_button.clicked.connect(self.inference_toggle)
        model_h_layout.addWidget(self.load_model_Button)
        model_h_layout.addWidget(self.infer_button)
        
        button_layout = QHBoxLayout()
        self.start_button = QPushButton('开启摄像头', self)
        self.start_button.clicked.connect(self.start_cam)
        self.stop_button = QPushButton('关闭摄像头', self)
        self.stop_button.clicked.connect(self.stop_cam)
        button_layout.addWidget(self.start_button)
        button_layout.addWidget(self.stop_button)
        
        p_h_layout = QHBoxLayout()
        self.p_value_label = QLabel(self)
        self.p_slider = QSlider(Qt.Horizontal, self)
        self.p_slider.setMinimum(0)
        self.p_slider.setMaximum(100)
        self.p_slider.setValue(self.yuntai.pan_pid.get_p() * 1000)
        self.p_value_label.setText("P: {:.3f}".format(self.yuntai.pan_pid.get_p()))
        self.p_slider.valueChanged.connect(self.p_value_changed)
        p_h_layout.addWidget(self.p_value_label)
        p_h_layout.addWidget(self.p_slider)

        i_h_layout = QHBoxLayout()
        self.i_value_label = QLabel(self)
        self.i_slider = QSlider(Qt.Horizontal, self)
        self.i_slider.setMinimum(0)
        self.i_slider.setMaximum(100)
        self.i_slider.setValue(self.yuntai.pan_pid.get_i() * 1000)
        self.i_value_label.setText("I: {:.3f}".format(self.yuntai.pan_pid.get_i()))
        self.i_slider.valueChanged.connect(self.i_value_changed)
        i_h_layout.addWidget(self.i_value_label)
        i_h_layout.addWidget(self.i_slider)

        d_h_layout = QHBoxLayout()
        self.d_value_label = QLabel(self)
        self.d_slider = QSlider(Qt.Horizontal, self)
        self.d_slider.setMinimum(0)
        self.d_slider.setMaximum(100)
        self.d_slider.setValue(self.yuntai.pan_pid.get_d() * 1000)
        self.d_value_label.setText("D: {:.3f}".format(self.yuntai.pan_pid.get_d()))
        self.d_slider.valueChanged.connect(self.d_value_changed)
        d_h_layout.addWidget(self.d_value_label)
        d_h_layout.addWidget(self.d_slider)
        
        sub_main_left_layout = QVBoxLayout()
        sub_main_left_layout.addWidget(self.image_label)
        sub_main_left_layout.addLayout(model_h_layout)
        sub_main_left_layout.addLayout(button_layout)
        sub_main_left_layout.addLayout(p_h_layout)
        sub_main_left_layout.addLayout(i_h_layout)
        sub_main_left_layout.addLayout(d_h_layout)
        
        self.sub_main_right_layout = QVBoxLayout()
        self.sub_main_right_layout.setAlignment(Qt.AlignTop)
        self.create_button_group()
        
        self.show_target_label = QLabel(self)
        self.show_target_label.setText("当前目标: None")
        self.show_target_label.setAlignment(Qt.AlignTop)
        self.show_target_label.setMinimumHeight(30)
        self.show_target_label.setMaximumHeight(30)
        
        explain = "使用说明:\n" \
                    "1. 加载模型\n" \
                    "2. 开启摄像头\n" \
                    "3. 开启推理\n" \
                    "4. 追踪目标\n" \
                    "  目标可为：pesron、flower pests None\n" \
                    "5. 关闭推理\n" \
                    "6. 关闭摄像头\n" \
                    "Time:  刘佳城  2025-4-15"
      
        explain_label = QLabel(explain, self)
        
        self.sub_main_right_layout.addWidget(self.show_target_label)
        self.sub_main_right_layout.addWidget(explain_label)
        
        self.main_layout = QHBoxLayout()
        self.main_layout.addLayout(sub_main_left_layout)
        self.main_layout.addLayout(self.sub_main_right_layout)
        self.setLayout(self.main_layout)

    def handle_debug_message(self, message):
        print(f"[DEBUG] {message}", flush=True)   

    def create_button_group(self):
        self.category_buttons = []
        top_row_layout = QHBoxLayout()
        bottom_row_layout = QHBoxLayout()

        for i, category in enumerate(categories):
            button = QPushButton(category, self)
            button.setFixedHeight(80)
            button.clicked.connect(lambda checked, category=category: self.set_target_name(category))
            self.category_buttons.append(button)
            if i < 2:
                top_row_layout.addWidget(button)
            else:
                bottom_row_layout.addWidget(button)

        button_layout = QVBoxLayout()
        button_layout.addLayout(top_row_layout)
        button_layout.addLayout(bottom_row_layout)
        self.sub_main_right_layout.addLayout(button_layout)
        self.video_thread.infer_results_signal.connect(self.update_button_colors)

    def update_image(self, cv_img):
        qt_img = QPixmap.fromImage(cv_img)
        self.image_label.setPixmap(qt_img)
        
    def start_cam(self):
        if not self.video_thread.isRunning():
            self.video_thread.start()
            
    def update_button_colors(self, infer_result):
        for button in self.category_buttons:
            button.setStyleSheet("")
        classid = infer_result["classid"]
        if classid is not None:
            for class_id in classid:
                self.category_buttons[int(class_id)].setStyleSheet("background-color: yellow")

    def set_target_name(self, target_name):
        target = target_name if self.yuntai_thread.target_name != target_name else None
        self.yuntai_thread.set_target_name(target)
        self.yolov11_wrapper.set_target_name(target)
        self.show_target_label.setText("当前目标：{}".format(target))

    def stop_cam(self):
        self.video_thread.stop()
        
    def load_model(self):
        if not self.yolov11_wrapper.load_model_status:
           self.yolov11_wrapper.load_model()
           self.load_model_status = True
           self.infer_button.setDisabled(False)
           self.load_model_Button.setDisabled(True)
        else:
            print("模型已经加载")
    
    def inference_toggle(self):
        self.video_thread.infer = not self.video_thread.infer
        self.infer_button.setText("关闭推理" if self.video_thread.infer else "开启推理")
    
    def p_value_changed(self, value):
        p = value / 1000
        self.yuntai.pan_pid.set_p(p)
        self.yuntai.tilt_pid.set_p(p)
        self.p_value_label.setText("P: {:.3f}".format(p))
        
    def i_value_changed(self, value):
        i = value / 1000
        self.yuntai.pan_pid.set_i(i)
        self.yuntai.tilt_pid.set_i(i)
        self.i_value_label.setText("I: {:.3f}".format(i))
        
    def d_value_changed(self, value):
        d = value / 1000
        self.yuntai.pan_pid.set_d(d)
        self.yuntai.tilt_pid.set_d(d)
        self.d_value_label.setText("D: {:.3f}".format(d))
        
    def closeEvent(self, event):
        print("应用程序正在关闭，执行清理工作...")
        self.video_thread.stop()
        time.sleep(1)
        yolov11_wrapper.destroy()
        self.yuntai_thread.stop()
        if cam.isOpened():
            cam.release()
        super().closeEvent(event)

if __name__ == "__main__":
    yolov11_wrapper = YoloV11TRT(Engine_file_path, Conf, Classid)
    device_id = 1
    cam = Camera(device_id)
    yuntai = YuntaiTrack()

    try:
        app = QApplication(sys.argv)
        a = App(yolov11_wrapper, cam, yuntai)
        a.show()
        sys.exit(app.exec_())
    except Exception as e:
        print(f"发生异常: {e}")
    finally:
        print("资源已释放")